import tushare as ts
from emmodels.em_base_data import BaseData, ModuleData
from emmodels.em_selffund import SelfFund
import datetime
import time
import emtencent.em_classified
import pandas as pd



def sql_base_data():
    now_str = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    t1 = ts.get_stock_basics()
    l1 = []
    t1_dict = t1.to_dict(orient='index')
    for code, data in t1_dict.items():
        ipodate = time.strptime(str(data['timeToMarket']), "%Y%m%d") if data['timeToMarket'] != 0 else None
        # base_data = BaseData(stkcode=code, stkname=data['name'], industry_zj=data['industry'],area=data['area'],ipodate=ipodate,create_time=now_str)
        data['code'] = code
        data['create_time'] = now_str
        data['ipodate'] = ipodate
        base_data = BaseData(**data)
        l1.append(base_data)
    BaseData.insert_list(l1, append_dup_list=['industry_zj', 'area', 'ipodate'])

    t2 = ts.get_concept_classified()
    l2 = []
    for name, data in t2.groupby(t2['code']):
        concept_sina = ','.join(data['c_name'].values)
        base_data = BaseData(stkcode=data['code'].values[0], stkname=data['name'].values[0], concept_sina=concept_sina)
        l2.append(base_data)
        # BaseData.insert(base_data,append_dup_list=['concept_sina'])
    BaseData.insert_list(l2, append_dup_list=['concept_sina'])

    tc_concept_classified = emtencent.em_classified.get_tc_concept_classified(pause=1)
    tc_concept_l = []
    for name, data in tc_concept_classified.groupby(tc_concept_classified['code']):
        concept_tc = ','.join(data['concept_name'].values)
        base_data = BaseData(stkcode=data['code'].values[0], stkname=data['name'].values[0], concept_tc=concept_tc)
        tc_concept_l.append(base_data)
    BaseData.insert_list(tc_concept_l, append_dup_list=['concept_tc'])

    sme = ts.get_sme_classified()
    l_sme = []
    data = {}
    for item in sme.values:
        base_data = BaseData(stkcode=item[0], stkname=item[1], sys_cls='中小板')
        # base_data = BaseData(**data)
        l_sme.append(base_data)
    BaseData.insert_list(l_sme, append_dup_list=['sys_cls'])

    sme = ts.get_gem_classified()
    l_sme = []
    for item in sme.values:
        base_data = BaseData(stkcode=item[0], stkname=item[1], sys_cls='创业板')
        l_sme.append(base_data)
    BaseData.insert_list(l_sme, append_dup_list=['sys_cls'])

    sme = ts.get_hs300s()
    l_sme = []
    if sme.values:
        for item in sme.values:
            base_data = BaseData(stkcode=item[0], stkname=item[1], sys_cls='沪深300')
            l_sme.append(base_data)
        BaseData.insert_list(l_sme, append_dup_list=['sys_cls'])

    sme = ts.get_sz50s()
    l_sme = []
    if sme.values:
        for item in sme.values:
            base_data = BaseData(stkcode=item[0], stkname=item[1], sys_cls='上证50')
            l_sme.append(base_data)
        BaseData.insert_list(l_sme, append_dup_list=['sys_cls'])

    now_str = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    sme = ts.get_zz500s()
    l_sme = []
    if sme.values:
        for item in sme.values:
            base_data = BaseData(stkcode=item[0], stkname=item[1], sys_cls='中证500', last_update_time=now_str)
            l_sme.append(base_data)
        BaseData.insert_list(l_sme, append_dup_list=['sys_cls'])

#基础数据
dic_base_data_cache = {}

# 实时数据，用下面这两个
dic_self_fund_data = {}
dic_module_data = {}

def mongo_base_data():
    global dic_base_data_cache
    now_str = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    t1 = ts.get_stock_basics()
    t1_dict = t1.to_dict(orient='index')
    for code, data in t1_dict.items():
        ipodate = datetime.datetime.strptime(str(data['timeToMarket']),'%Y%m%d') if data['timeToMarket'] != 0 else None
        data['code'] = code
        data['create_time'] = now_str
        data['ipodate'] = ipodate
        base_data = BaseData(**data)
        dic_base_data_cache[code] = base_data

    try:
        t2 = ts.get_concept_classified()
        for name, data in t2.groupby(t2['code']):
            concept_sina = ','.join(data['c_name'].values)
            code = data['code'].values[0]
            if code in dic_base_data_cache.keys():
                base_data = dic_base_data_cache[code]
                base_data.concept_sina = concept_sina
    except Exception as e:
        print(e)



    tc_concept_classified = emtencent.em_classified.get_tc_concept_classified(pause=1)
    for name, data in tc_concept_classified.groupby(tc_concept_classified['code']):
        concept_tc = ','.join(data['concept_name'].values)
        code = data['code'].values[0]
        if code in dic_base_data_cache.keys():
            base_data = dic_base_data_cache[code]
            base_data.concept_tc = concept_tc

    tc_inductry_classified = emtencent.em_classified.get_tc_industry_classified(pause=0.5)
    for name, data in tc_inductry_classified.groupby(tc_inductry_classified['code']):
        inductry_tc = ','.join(data['industry_name'].values)
        code = data['code'].values[0]
        if code in dic_base_data_cache.keys():
            base_data = dic_base_data_cache[code]
            base_data.industry_zj = inductry_tc

    pass
    # sme = ts.get_sme_classified()
    # if not sme.empty:
    #     for item in sme.values:
    #         code = item[0]
    #         if code in dic_base_data_cache.keys():
    #             base_data = dic_base_data_cache[code]
    #             base_data.sys_cls = '中小板'
    #
    # sme = ts.get_gem_classified()
    # if not sme.empty:
    #     for item in sme.values:
    #         code = item[0]
    #         if code in dic_base_data_cache.keys():
    #             base_data = dic_base_data_cache[code]
    #             base_data.sys_cls = '创业板'
    #
    #
    #
    # sme = ts.get_hs300s()
    # if not sme.empty:
    #     for item in sme.values:
    #         code = item[0]
    #         if code in dic_base_data_cache.keys():
    #             base_data = dic_base_data_cache[code]
    #             base_data.sys_cls = '沪深300'
    #
    # sme = ts.get_sz50s()
    # if not sme.empty:
    #     for item in sme.values:
    #         code = item[0]
    #         if code in dic_base_data_cache.keys():
    #             base_data = dic_base_data_cache[code]
    #             base_data.sys_cls = '上证50'
    #
    # now_str = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    # sme = ts.get_zz500s()
    # if not sme.empty:
    #     for item in sme.values:
    #         code = item[0]
    #         if code in dic_base_data_cache.keys():
    #             base_data = dic_base_data_cache[code]
    #             base_data.sys_cls = '中证500'
    #             base_data.last_update_time = now_str
    #
    for key, item in dic_base_data_cache.items():
        item.saveToMongo()

def mongo_module_data():

    tc_concept_classified = emtencent.em_classified.get_tc_concept_classified(pause=1)
    for name, data in tc_concept_classified.groupby(tc_concept_classified['concept_code']):
        stk_codes = data['code'].values
        stk_names = data['name'].values
        infos = stk_names+'('+stk_codes+')'
        moduleData = ModuleData(module_code=data['concept_code'].values[0], module_name=data['concept_name'].values[0],
                                module_stk_codes=list(stk_codes), module_stk_names=list(stk_names),
                                module_stk_infos=list(infos), source='tc_concept')
        moduleData.saveToMongo()

    tc_inductry_classified = emtencent.em_classified.get_tc_industry_classified(pause=0.5)
    for name, data in tc_inductry_classified.groupby(tc_inductry_classified['industry_code']):
        stk_codes = data['code'].values
        stk_names = data['name'].values
        infos = stk_names+'('+stk_codes+')'
        moduleData = ModuleData(module_code=data['industry_code'].values[0], module_name=data['industry_name'].values[0],
                                module_stk_codes=list(stk_codes), module_stk_names=list(stk_names),
                                module_stk_infos=list(infos), source='tc_industry')
        moduleData.saveToMongo()


def cache_datas_from_mongo():
    get_base_from_mongo()
    get_module_from_mongo()

def get_base_from_mongo():
    # 缓存基础数据
    if not len(dic_base_data_cache):
        r = BaseData.select_all_from_mongo()
        df = pd.DataFrame(data=list(r), columns=BaseData._db_mapper.values())
        # df = df.set_index(df['stkcode'])
        for key, value in df.iterrows():
            dic = value.to_dict()
            baseData = BaseData(**dic)
            dic_base_data_cache[dic['stkcode']] = baseData
        return df

def get_module_from_mongo():
    # 缓存模块数据
    if not len(dic_module_data):
        l = list(ModuleData.select_all_from_mongo())
        df = pd.DataFrame(data=l)
        for value in l:
            moduleData = ModuleData(**value)
            dic_module_data[moduleData.module_code] = moduleData
        return df




def get_self_fund(stkcode):
    if not len(dic_self_fund_data) and not len(dic_base_data_cache):
        return None
    if stkcode in dic_self_fund_data.keys():
        return dic_self_fund_data[stkcode]
    else:
        sfund = SelfFund.select_stkcode(stkcode)
        dic_self_fund_data[stkcode] = sfund
        return sfund

def get_module(module_code):
    if not len(dic_module_data):
        return None
    if module_code in dic_module_data.keys():
        return dic_module_data[module_code]
    else:
        r = ModuleData.select_module(module_code)
        if r:
            moduleData = ModuleData(**r)
            dic_module_data[module_code] = moduleData
            return moduleData
        return None



cache_datas_from_mongo()

if __name__ == '__main__':
    # get_self_fund('002241')
    mongo_base_data()
    # mongo_module_data()
    # cache_datas_from_mongo()
    pass

